# 计算插值的系数是非常昂贵的, 所以我们需要把系数存下来.
# 为了方便起见, 我们针对 n 等分插值点直接特化.

import numpy as np
from scipy.linalg import solve_banded  # 用来解转角方程


class Spline:
    def __init__(self, num: int):
        """ 在 [-1,1] 上的 num 等分点插值 """
        assert (num >= 2)
        self.n = num
        self._compute_coefficients()

    def compute(self, x: float):
        """ 对 x, 给出三次 spline 插值的结果 """
        assert (-1 <= x <= 1)
        i = int((x + 1) * self.n / 2)
        i = i if i < self.n else self.n - 1
        s = self.f[i] * alpha(x, i, self.n) + self.m[i] * beta(x, i, self.n) + self.f[i] * alpha(x, i + 1, self.n) + \
            self.m[i + 1] * beta(x, i, self.n)
        return s

    def compute_array(self, arr):
        """ 对数组 arr, 给出三次 spline 插值的结果 """
        y = []
        for x in arr:
            y.append(self.compute(x))
        return y

    def _compute_coefficients(self):
        """ 计算三次插值多项式的系数 """
        self.f = np.zeros(self.n + 1)
        for j in range(self.n + 1):
            self.f[j] = f(-1 + 2 * j / self.n)

        A = np.zeros((3, self.n + 1))
        # 主对角线
        A[1, :] = 2
        # 上方的对角线
        A[0, 2:] = 1 / 2
        A[0, 1] = 1
        # 下方的对角线
        A[-1, :-2] = 1 / 2
        A[-1, -2] = 1

        b = np.empty((self.n + 1, 1), dtype=float)
        b[0] = 3 * df(0, self.n)
        b[-1] = 3 * df(self.n - 1, self.n)
        for j in range(1, self.n):
            b[j] = 3 * (df(j - 1, self.n) + df(j, self.n)) / 2

        # 解转角方程
        self.m = solve_banded((1, 1), A, b, overwrite_ab=True, overwrite_b=True, check_finite=False).reshape(self.n + 1)


def f(x: float):
    return 1 / (1 + 25 * x ** 2)


def df(j: int, n: int):
    """ 返回 f[x_j, x_{j+1}] """
    assert (0 <= j <= n - 1)
    return (f(-1 + 2 * (j + 1) / n) - f(-1 + 2 * j / n)) * n / 2


def alpha(x: float, j: int, n: int):
    xj = -1 + 2 * j / n
    if (- 2 / n) <= (x - xj) <= 0:
        return (1 - n * x - n + 2 * j) * (n * x / 2 - (j - 1) + n / 2) ** 2
    elif 0 <= (x - xj) <= (2 / n):
        return (n * x + n + 1 - 2 * j) * (n * x / 2 - (j + 1) + n / 2) ** 2
    else:
        return 0


def beta(x: float, j: int, n: int):
    xj = -1 + 2 * j / n
    if (- 2 / n) <= (x - xj) <= 0:
        return (x - xj) * (n * x / 2 - (j - 1) + n / 2) ** 2
    elif 0 <= (x - xj) <= (2 / n):
        return (x - xj) * (n * x / 2 - (j + 1) + n / 2) ** 2
    else:
        return 0
